"""
机器学习策略(随机森林分类)(Zipline实现)
策略逻辑：
1. 使用技术指标作为特征
2. 预测未来N日涨跌概率
3. 当上涨概率高时买入，下跌概率高时卖出
"""

from zipline.api import order_target_percent, record, symbol, get_datetime
from zipline.finance import commission, slippage
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier

def initialize(context):
    # 策略参数
    context.asset = symbol('SPY')
    context.lookback = 30  # 特征计算周期
    context.hold_period = 5  # 持仓周期
    context.n_estimators = 100  # 随机森林树数量
    context.prob_threshold = 0.7  # 交易概率阈值
    
    # 设置交易成本
    context.set_commission(commission.PerShare(cost=0.001, min_trade_cost=1))
    context.set_slippage(slippage.FixedSlippage(spread=0.01))
    
    # 初始化模型
    context.model = RandomForestClassifier(
        n_estimators=context.n_estimators,
        random_state=42
    )
    
    # 初始化数据缓存
    context.X = []  # 特征数据
    context.y = []  # 标签数据
    context.feature_history = []  # 特征历史
    context.price_history = []  # 价格历史

def calculate_features(data, context):
    """计算技术指标特征"""
    # 获取历史价格
    prices = np.array(context.price_history)
    
    # 计算特征
    pct_change = (prices[-1] - prices[-2]) / prices[-2] if len(prices) > 1 else 0
    
    ma5 = np.mean(prices[-5:]) if len(prices) >= 5 else prices[-1]
    ma20 = np.mean(prices[-20:]) if len(prices) >= 20 else prices[-1]
    ma_diff = (ma5 - ma20) / ma20
    
    # RSI计算
    if len(prices) >= 15:
        gains = np.where(np.diff(prices) > 0, np.diff(prices), 0)[-14:]
        losses = np.where(np.diff(prices) < 0, -np.diff(prices), 0)[-14:]
        avg_gain = np.mean(gains)
        avg_loss = np.mean(losses)
        rsi = 100 - (100 / (1 + (avg_gain / avg_loss if avg_loss != 0 else 1)))
    else:
        rsi = 50
    
    # MACD计算
    if len(prices) >= 26:
        ema12 = pd.Series(prices).ewm(span=12).mean().values[-1]
        ema26 = pd.Series(prices).ewm(span=26).mean().values[-1]
        macd = ema12 - ema26
    else:
        macd = 0
    
    return [pct_change, ma_diff, rsi, macd]

def handle_data(context, data):
    # 获取当前价格
    current_price = data.current(context.asset, 'price')
    context.price_history.append(current_price)
    
    # 保持价格历史长度
    if len(context.price_history) > context.lookback + context.hold_period:
        context.price_history.pop(0)
    
    # 检查是否有足够数据
    if len(context.price_history) < context.lookback + context.hold_period:
        return
    
    # 计算特征
    features = calculate_features(data, context)
    context.feature_history.append(features)
    
    # 添加训练数据 (每20天一次)
    if len(context.feature_history) % 20 == 0:
        # 获取标签 (未来hold_period天的收益)
        future_price = context.price_history[-context.hold_period]
        future_return = (future_price - current_price) / current_price
        label = 1 if future_return > 0 else 0
        
        context.X.append(features)
        context.y.append(label)
    
    # 定期重新训练模型 (当有足够样本)
    if len(context.X) >= 100 and len(context.X) % 20 == 0:
        context.model.fit(context.X, context.y)
    
    # 预测当前市场
    if hasattr(context.model, 'predict_proba') and len(context.X) >= 100:
        proba = context.model.predict_proba([features])[0][1]
        
        # 生成交易信号
        if proba > context.prob_threshold:
            # 上涨概率高，买入
            order_target_percent(context.asset, 1.0)
        elif proba < (1 - context.prob_threshold):
            # 下跌概率高，卖出
            order_target_percent(context.asset, -1.0)
        else:
            # 中性，保持仓位
            pass
        
        # 记录状态
        record(price=current_price, 
              proba=proba, 
              position=context.portfolio.positions[context.asset].amount)